package com.xagent.engine.context;

import com.xagent.engine.context.inter.InstWrapper;
import com.xagent.engine.utils.HookUtil;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.Instrumentation;
import java.lang.instrument.UnmodifiableClassException;
import java.util.*;

/**
 * Instrumentation的Wrapper类
 * 用以提供instrumentation并且封装一些常用功能
 */
public class InstWrapperImpl implements InstWrapper {
    private Instrumentation instrumentation;
    private List<ClassFileTransformer> transformerList;
    private final Logger agentLogger = LoggerFactory.getLogger("agentLogger");
    private final Logger securityLogger = LoggerFactory.getLogger("securityLogger");


    private final Set<ClassLoader> classLoaders = new HashSet<ClassLoader>();

    public InstWrapperImpl(Instrumentation instrumentation) {
        this.instrumentation = instrumentation;
        transformerList = new ArrayList<ClassFileTransformer>();
    }

    /**
     * 将transformer添加进instrumentation，并且记录在transformerList中
     */
    @Override
    public void addTransformer(ClassFileTransformer transformer) {

        // 不重复添加
        if (transformer == null || transformerList.contains(transformer)) return;
        instrumentation.addTransformer(transformer, true);
    }

    @Override
    public void retransformClasses(Class<?> clazz) throws UnmodifiableClassException {
        if (clazz == null) return;
        instrumentation.retransformClasses(clazz);
    }

    @Override
    public void retransformAllClasses() throws UnmodifiableClassException {
        HashSet<Class<?>> classes = new HashSet<>();
        Class[] allLoadedClasses = instrumentation.getAllLoadedClasses();

        for (Class aClazz : allLoadedClasses) {
            classes.add(aClazz);
            addClassLoader(aClazz.getClassLoader());
        }

        // 这里需要考虑是直接在上一个for循环中直接进行retransform好还是新起一个循环进行retransform好，暂时先新起循环
        for (Class<?> aClass : classes) {
            if (AgentContext.getContext().getInstWrapper().filterOrEnhance(aClass.getName())) continue;

            retransformClasses(aClass);
        }

    }

    @Override
    public Instrumentation getInstrumentation() {
        return instrumentation;
    }

    /**
     * 删除已添加的指定transformer
     */
    @Override
    public void removeTransformer(ClassFileTransformer transformer) {
        if (transformer == null) return;
        if (transformerList.contains(transformer)) {
            transformerList.remove(transformer);
            instrumentation.removeTransformer(transformer);
        }
    }

    @Override
    public boolean addClassLoader(ClassLoader classLoader) {
        if (classLoader == null) return false;

        if (classLoaders.add(classLoader)) return true;

        return false;
    }

    /**
     * @param className 待检测的类名
     * @return true 代表不hook，false代表需要hook.
     */
    @Override
    public boolean filterOrEnhance(String className) {
        /**
         * 不需要Hook的包
         * */
        for (String excludePackage : HookUtil.excludePackages) {
            if (className.startsWith(excludePackage)) return true;
        }

        /**
         * 异常类不需要Hook
         * */
        if (className.endsWith("Exception")) return true;

        return false;
    }
}
